#!/usr/bin/env python3
"""
音频收集系统 - 句子分配功能数据库迁移脚本
用于在现有系统上添加句子分配管理功能

使用方法:
python migrate_assignment_tables.py

注意：运行前请确保已备份数据库
"""

import sys
import os
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from app import app
from models.models import db, User, Sentence, Recording, UserSentenceAssignment, AssignmentTemplate

def check_tables_exist():
    """检查新表是否已存在"""
    try:
        with app.app_context():
            # 尝试查询新表
            UserSentenceAssignment.query.first()
            AssignmentTemplate.query.first()
            print("✅ 分配表已存在，跳过创建")
            return True
    except Exception:
        print("📋 分配表不存在，需要创建")
        return False

def create_tables():
    """创建新的数据库表"""
    try:
        with app.app_context():
            print("🔨 正在创建分配管理相关表...")
            db.create_all()
            print("✅ 表创建成功")
            return True
    except Exception as e:
        print(f"❌ 表创建失败: {e}")
        return False

def migrate_existing_recordings():
    """为现有录音创建分配记录"""
    try:
        with app.app_context():
            print("🔄 正在迁移现有录音数据...")
            
            # 获取所有录音记录
            recordings = Recording.query.all()
            print(f"📊 发现 {len(recordings)} 条录音记录")
            
            migrated_count = 0
            skipped_count = 0
            
            # 按用户分组处理
            user_recordings = {}
            for recording in recordings:
                if recording.user_id not in user_recordings:
                    user_recordings[recording.user_id] = []
                user_recordings[recording.user_id].append(recording)
            
            for user_id, user_recordings_list in user_recordings.items():
                # 按创建时间排序
                user_recordings_list.sort(key=lambda x: x.id)
                
                for i, recording in enumerate(user_recordings_list, 1):
                    # 检查是否已存在分配记录
                    existing = UserSentenceAssignment.query.filter_by(
                        user_id=recording.user_id,
                        sentence_id=recording.sentence_id
                    ).first()
                    
                    if existing:
                        skipped_count += 1
                        continue
                    
                    # 创建分配记录
                    assignment = UserSentenceAssignment(
                        user_id=recording.user_id,
                        sentence_id=recording.sentence_id,
                        sequence_order=i,
                        status='completed',
                        assigned_at=datetime.utcnow(),
                        completed_at=datetime.utcnow()
                    )
                    
                    db.session.add(assignment)
                    migrated_count += 1
            
            db.session.commit()
            print(f"✅ 迁移完成: 新增 {migrated_count} 条分配记录，跳过 {skipped_count} 条已存在记录")
            return True
            
    except Exception as e:
        db.session.rollback()
        print(f"❌ 数据迁移失败: {e}")
        return False

def create_sample_templates():
    """创建示例分配模板"""
    try:
        with app.app_context():
            print("📝 正在创建示例分配模板...")
            
            # 检查是否已有模板
            existing_templates = AssignmentTemplate.query.count()
            if existing_templates > 0:
                print(f"📋 已存在 {existing_templates} 个模板，跳过创建示例模板")
                return True
            
            # 获取前15个句子用于创建模板
            sentences = Sentence.query.limit(15).all()
            if len(sentences) < 15:
                print(f"⚠️  句子数量不足（只有{len(sentences)}个），跳过创建示例模板")
                return True
            
            # 创建示例模板
            templates = [
                {
                    'name': '新疆风光5句套餐',
                    'description': '包含新疆风光相关的5个句子，适合初次录制',
                    'sentence_ids': [s.id for s in sentences[:5]],
                    'category': '风光类'
                },
                {
                    'name': '基础普通话测试',
                    'description': '基础普通话水平测试句子',
                    'sentence_ids': [s.id for s in sentences[5:10]],
                    'category': '测试类'
                },
                {
                    'name': '日常对话练习',
                    'description': '日常生活对话相关句子',
                    'sentence_ids': [s.id for s in sentences[10:15]],
                    'category': '对话类'
                }
            ]
            
            created_count = 0
            for template_data in templates:
                template = AssignmentTemplate(**template_data)
                db.session.add(template)
                created_count += 1
            
            db.session.commit()
            print(f"✅ 创建了 {created_count} 个示例模板")
            return True
            
    except Exception as e:
        db.session.rollback()
        print(f"❌ 创建示例模板失败: {e}")
        return False

def print_migration_summary():
    """打印迁移总结"""
    try:
        with app.app_context():
            print("\n" + "="*50)
            print("📊 迁移总结")
            print("="*50)
            
            # 统计信息
            total_users = User.query.filter_by(role='user').count()
            total_sentences = Sentence.query.count()
            total_recordings = Recording.query.count()
            total_assignments = UserSentenceAssignment.query.count()
            total_templates = AssignmentTemplate.query.count()
            
            # 有分配的用户数
            users_with_assignments = db.session.query(UserSentenceAssignment.user_id).distinct().count()
            
            print(f"👥 总用户数: {total_users}")
            print(f"📝 总句子数: {total_sentences}")
            print(f"🎵 总录音数: {total_recordings}")
            print(f"📋 总分配数: {total_assignments}")
            print(f"📄 总模板数: {total_templates}")
            print(f"👤 有分配的用户: {users_with_assignments}")
            
            print("\n🎉 句子分配功能迁移完成！")
            print("\n📖 新增API接口:")
            print("   管理端:")
            print("   - GET  /api/admin/assignments/overview")
            print("   - GET  /api/admin/assignments/user/<user_id>")
            print("   - POST /api/admin/assignments/assign")
            print("   - POST /api/admin/assignments/assign-random")
            print("   - POST /api/admin/assignments/batch-assign")
            print("   用户端:")
            print("   - GET  /api/sentences/assigned")
            print("   - GET  /api/sentences/next")
            print("   - GET  /api/sentences/progress")
            
    except Exception as e:
        print(f"❌ 获取统计信息失败: {e}")

def main():
    """主函数"""
    print("🚀 开始句子分配功能数据库迁移")
    print("="*50)
    
    # 检查表是否存在
    if check_tables_exist():
        print("⚠️  表已存在，只进行数据迁移")
    else:
        # 创建表
        if not create_tables():
            print("❌ 迁移失败：无法创建表")
            return False
    
    # 迁移现有数据
    if not migrate_existing_recordings():
        print("❌ 迁移失败：数据迁移出错")
        return False
    
    # 创建示例模板
    if not create_sample_templates():
        print("⚠️  示例模板创建失败，但不影响主要功能")
    
    # 打印总结
    print_migration_summary()
    
    return True

if __name__ == '__main__':
    try:
        success = main()
        if success:
            print("\n✅ 迁移成功完成！")
            sys.exit(0)
        else:
            print("\n❌ 迁移失败！")
            sys.exit(1)
    except KeyboardInterrupt:
        print("\n⚠️  迁移被用户中断")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 迁移过程中发生未知错误: {e}")
        sys.exit(1)
